import numpy as np
import torch
from StreamLearn.Base.SemiEstimator import StreamEstimator


class GDRO(StreamEstimator):
    def __init__(self, args):
        self.name = 'GDRO'
        self.D = args.D
        self.G = args.G
        self.cw = args.cw_GDRO
        self.cq = args.cq_GDRO
        self.T = None
        self.m = None
        self.w_size = None
        # Step_sizes
        self.eta_w = None
        self.eta_q = None
        # Parameters
        self.w = None
        self.ave_w = None
        self.q = None
        self.ave_q = None
        # Optimization setting
        self.criterion = torch.nn.BCELoss(reduction='none')
        self.function = torch.nn.Sigmoid()
        self.optimizer = None

    def init(self):
        # Step_sizes
        self.eta_w = torch.tensor(self.cw * self.D ** 2 * np.sqrt(1 / (self.T * (self.D ** 2 * self.G ** 2 +
                                                                                 np.log(self.m)))))
        self.eta_q = torch.tensor(self.cq * np.log(self.m) * np.sqrt(1 / (self.T * (self.D ** 2 * self.G ** 2 +
                                                                                    np.log(self.m)))))
        # Parameters
        self.w = torch.zeros(self.w_size, requires_grad=True, dtype=torch.float32)
        self.ave_w = torch.zeros(self.w_size)
        self.q = torch.ones(self.m) / self.m
        self.ave_q = torch.zeros(self.m)
        # Optimization setting
        self.optimizer = torch.optim.SGD([self.w], lr=self.eta_w.numpy())

    def fit(self, stream_dataset):
        self.m = stream_dataset.get_m()
        self.w_size = stream_dataset.get_w_size()
        self.T = np.floor(np.min(stream_dataset.get_n())).astype(int)
        self.init()
        for t in range(1, self.T + 1):
            self.ave_q = (self.ave_q * (t - 1) + self.q) / t
            self.ave_w = (self.ave_w * (t - 1) + self.w.detach()) / t
            sample_m = stream_dataset.sample_m()
            xs = sample_m[:, :-1]
            ys = sample_m[:, -1]
            model_output = self.function(torch.from_numpy(xs).to(torch.float32) @ self.w)
            loss = self.criterion(model_output, torch.tensor(ys, dtype=torch.float32))
            true_loss = self.q @ loss
            self.optimizer.zero_grad()
            true_loss.backward()
            self.optimizer.step()
            if torch.norm(self.w.detach(), p=2) > self.D:
                self.w = self.w.detach() / torch.norm(self.w.detach(), p=2) * self.D
                self.w.requires_grad_(True)
            q_1 = self.q * np.exp(self.eta_q * loss.detach().numpy())
            q_1 = q_1 / q_1.sum()
            self.q = q_1

    def predict(self, X):
        model_output = self.function(torch.from_numpy(X).to(torch.float32) @ self.ave_w)
        return model_output

    def evaluate(self, y_pred, y_true):
        return self.criterion(y_pred, torch.tensor(y_true, dtype=torch.float32)).numpy()
